#!/usr/bin/python
# -*- coding:utf-8 -*-

import datetime
from pymongo import MongoClient

from conf.settings import (
    CLOUDMONITOR_MONGODB,
    CLOUDMONITOR_DB,
    IGNORE_TYPE,
    MONITOR_TIME,
    TASK_ERROR_TIME
)
from utility.wx_notify import notify
from utility.log import Logger

logger = Logger(filename='monitor_tasks.log').get_logger()


db = MongoClient(CLOUDMONITOR_MONGODB)[CLOUDMONITOR_DB]
# 数据存储时间为utc时间
now = datetime.datetime.now() - datetime.timedelta(hours=8)
last_time = now - datetime.timedelta(hours=MONITOR_TIME)

task_map = {
    'vul': {"name": "WEB漏洞", "time": TASK_ERROR_TIME.get("vul", 1800), "count": 0, "task_id": [], "taskSessionId": []},
    'ssl': {"name": "SSL安全", "time": TASK_ERROR_TIME.get("ssl", 420), "count": 0, "task_id": [], "taskSessionId": []},
    'asset': {"name": "资产变动", "time": TASK_ERROR_TIME.get("asset", 1800), "count": 0, "task_id": [], "taskSessionId": []},
    'securityEvent': {"name": "安全事件", "time": TASK_ERROR_TIME.get("securityEvent", 360), "count": 0, "task_id": [], "taskSessionId": []},
    'content': {"name": "内容合规", "time": TASK_ERROR_TIME.get("content", 240), "count": 0, "task_id": [], "taskSessionId": []},
    'sitePortrait': {"name": "网站画像", "time": TASK_ERROR_TIME.get("sitePortrait", 60), "count": 0, "task_id": [], "taskSessionId": []},
    'ping': {"name": "ping", "time": TASK_ERROR_TIME.get("ping", 400), "count": 0, "task_id": [], "taskSessionId": []},
    'http': {"name": "http", "time": TASK_ERROR_TIME.get("http", 400), "count": 0, "task_id": [], "taskSessionId": []},
}


def monitor_task_status():
    """监测任务执行状态，长时间没有执行完成"""
    query = {
        "endTime": {"$gt": last_time},
        "status": {"$in": ['active', 'waiting']}
    }
    display = {"startTime": 1, "taskType": 1, "_id": 1, "taskSessionId": 1}
    cursor = db.tasks.find(query, display)
    for task in cursor:
        start_time = task["startTime"]
        task_type = task["taskType"]
        if task_type in IGNORE_TYPE:
            continue
        if (now - start_time).total_seconds() >= task_map[task_type]["time"]:
            task_map[task_type]["count"] += 1
            task_map[task_type]["task_id"].append(task["_id"])
            task_map[task_type]["taskSessionId"].append(task.get("taskSessionId", ""))
    return task_map


def main():
    message = []
    monitor_task_status()
    for t, v in task_map.items():
        if v["count"] > 0:
            log_msg = f"有{v['count']}条{t}({v['name']})任务可能异常, taskSessionId: " \
                      f"{v['taskSessionId']}, task_id: {v['task_id']};"
            logger.info(log_msg)
        if (t in ["http", "ping"] and v['count'] > 20) or (t not in ["http", "ping"] and v['count'] > 10):
            id_str = f"taskSessionId有{v['taskSessionId'][:3]}" if v["taskSessionId"] else f"task_id有{v['task_id'][:3]}"
            message.append(f"有{v['count']}条{t}({v['name']})任务可能异常, {id_str};")
    if message:
        notify("\n".join(message))


if __name__ == "__main__":
    main()
